import torch
from torch import nn
import triton
import triton.language as tl

from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
from nanovllm.utils.context import get_context

"""
这段代码实现了一个高性能注意力层（Attention），核心整合了 Triton 内核优化（用于高效 KV 缓存存储）和 FlashAttention 加速库（用于快速注意力计算），专门适配大语言模型（LLM）的推理场景（支持 prefill 预填充和 decode 解码两个阶段）。以下从模块拆分、核心逻辑、优化细节展开分析：
一、核心依赖与背景
代码依赖两个关键加速工具，需先明确其作用：

Triton：NVIDIA 开源的 GPU 编程框架，可通过 Python 级别的内核编写，实现比 PyTorch 原生操作更高效的 GPU 内存访问与计算（减少 Python 层开销）。
FlashAttention：斯坦福开源的注意力计算库，通过算子融合和内存高效访问（避免重复加载 KV 数据），大幅降低注意力计算的时间与显存占用（比传统注意力快 2-4 倍）。
二、模块拆分与功能解析
代码分为 3 个核心部分：Triton KV 缓存存储内核、KV 缓存存储函数、Attention 类（含 prefill/decode 逻辑）。

核心逻辑拆解：
QKV 重塑：将输入的 “展平特征”（num_heads*head_dim）拆分为 “多头结构”（num_heads, head_dim），符合注意力计算的格式要求。
上下文获取：get_context() 是自定义工具，存储推理的全局信息（如是否为 prefill、KV 缓存的 slot 映射、块表等），避免参数冗余传递。
KV 缓存写入：若缓存已初始化，调用 store_kvcache 用 Triton 内核将当前 KV 写入缓存，供后续 decode 阶段复用。
Prefill 阶段：
用 flash_attn_varlen_func：支持变长序列批量处理（如同时处理多个不同长度的用户提问），通过 cu_seqlens_q/cu_seqlens_k 记录每个序列的长度边界。
前缀缓存（prefix cache）：若 block_tables 存在，直接复用缓存的 KV（如重复的提示词前缀，避免重复计算）。
Decode 阶段：
用 flash_attn_with_kvcache：专门为 “单 token 生成” 优化，仅计算新 Q 与历史 KV 的注意力，复用缓存，大幅减少计算量。
q.unsqueeze(1)：增加 batch 维度，适配 FlashAttention 接口要求。
输出重塑：将多头结果重新展平为 “单特征维度”，与输入格式对齐，确保后续网络层（如 Feed-Forward）可正常处理。
三、关键优化点总结
该 Attention 层是工业级 LLM 推理的典型实现，核心优化围绕 “速度” 和 “显存” 展开：

FlashAttention 加速：
避免 KV 数据的重复加载（传统注意力需多次加载 KV 到 GPU 寄存器），通过算子融合将注意力计算的多个步骤（得分计算、softmax、加权求和）合并为单个算子，减少内存带宽消耗。
支持因果掩码（causal=True），符合 LLM 生成式任务的逻辑（不能看到未来 token）。
Triton 内核优化 KV 存储：
直接操作 GPU 内存，跳过 PyTorch 张量的 Python 层封装，减少 overhead。
批量内存访问模式，适配 GPU 并行架构，提升 KV 缓存写入速度。
双阶段适配：
Prefill 阶段：处理长序列批量输入，用变长接口支持不同长度序列。
Decode 阶段：复用 KV 缓存，仅计算新 token 的注意力，推理速度提升 10-100 倍（相比无缓存的暴力计算）。
GQA/MQA 支持：
通过 num_kv_heads < num_heads 减少 KV 的数量，显存占用降低（如 num_kv_heads=8、num_heads=32，KV 显存减少 75%），同时保持模型表达能力。
四、应用场景
该 Attention 层是 LLM 推理引擎（如 nano-vllm、vllm）的核心组件，适用于：

大模型在线推理（如对话机器人、文本生成 API）：需低延迟、高并发处理用户请求。
长序列生成任务（如小说续写、代码生成）：依赖 KV 缓存复用，避免长序列重复计算。
资源受限场景（如消费级 GPU）：通过 GQA/MQA 和缓存优化，降低显存占用。
总结
这段代码是高性能 LLM 注意力层的工程化实现，通过 Triton 优化 KV 缓存存储、FlashAttention 加速注意力计算，同时适配 prefill/decode 双阶段推理，平衡了速度、显存与通用性。理解其逻辑后，可清晰掌握工业级 LLM 推理中 “如何通过底层优化提升性能” 的核心思路。
"""


@triton.jit  # Triton 即时编译装饰器：将函数编译为 GPU 内核
def store_kvcache_kernel(
        key_ptr,  # K 张量的 GPU 内存指针
        key_stride,  # K 张量第 0 维度的步长（用于计算元素偏移）
        value_ptr,  # V 张量的 GPU 内存指针
        value_stride,  # V 张量第 0 维度的步长
        k_cache_ptr,  # K 缓存的 GPU 内存指针
        v_cache_ptr,  # V 缓存的 GPU 内存指针
        slot_mapping_ptr,  # slot 映射表指针（记录当前 K/V 应写入缓存的位置）
        D: tl.constexpr,  # 单个 K/V 向量的维度（num_kv_heads * head_dim，编译期常量）
):
    """
    核心作用：将当前计算的 K（Key）、V（Value）张量高效写入 KV 缓存（推理中复用 KV 避免重复计算）。
优化点：
编译期常量 D：tl.constexpr 标记 D 为编译期已知，Triton 可提前优化内存访问模式。
批量偏移计算：用 tl.arange(0, D) 一次性计算维度内所有偏移，避免循环，提升并行效率。
直接 GPU 指针操作：跳过 PyTorch 张量的封装层，直接访问 GPU 内存，减少 overhead。
    """
    # 1. 每个线程处理一个 K/V 元素（按 program_id 分配索引）
    idx = tl.program_id(0)  # 线程索引（对应 K/V 张量的第 0 维度，如 [N, ...] 中的 N）

    # 2. 计算 K/V 张量中当前元素的内存偏移（按步长 + 维度遍历）
    key_offsets = idx * key_stride + tl.arange(0, D)  # K 元素的偏移：idx 对应第 N 个样本，遍历 D 维
    value_offsets = idx * value_stride + tl.arange(0, D)  # V 元素的偏移

    # 3. 从 K/V 张量加载数据到 GPU 寄存器
    key = tl.load(key_ptr + key_offsets)
    value = tl.load(value_ptr + value_offsets)

    # 4. 获取当前元素应写入缓存的 slot（位置）
    slot = tl.load(slot_mapping_ptr + idx)  # slot_mapping 记录 N 个元素对应的缓存位置

    # 5. 计算缓存中的偏移，将 K/V 写入缓存
    cache_offsets = slot * D + tl.arange(0, D)  # 缓存位置 = slot * 维度 + 维度内索引
    tl.store(k_cache_ptr + cache_offsets, key)  # 写入 K 缓存
    tl.store(v_cache_ptr + cache_offsets, value)  # 写入 V 缓存


def store_kvcache(
        key: torch.Tensor,  # 输入 K 张量（形状：[N, num_kv_heads, head_dim]）
        value: torch.Tensor,  # 输入 V 张量（形状：[N, num_kv_heads, head_dim]）
        k_cache: torch.Tensor,  # K 缓存张量（预分配的大缓存）
        v_cache: torch.Tensor,  # V 缓存张量（预分配的大缓存）
        slot_mapping: torch.Tensor  # slot 映射表（形状：[N]，记录每个 K/V 对应的缓存位置）
):
    """
    核心作用：封装 Triton 内核调用，确保输入张量满足内存连续性要求（避免 Triton 访问混乱），并传递必要参数。
步长校验的意义：Triton 内核依赖 “连续内存访问” 提升效率，若张量内存不连续（如切片后未 contiguous()），会导致访问错误或性能下降，因此通过 assert 提前规避。
    """
    # 1. 解析张量形状：N（元素数）、num_kv_heads（KV 头数）、head_dim（头维度）
    N, num_heads, head_dim = key.shape
    D = num_heads * head_dim  # 单个 K/V 元素的总维度（展平后的长度）

    # 2. 校验张量步长（确保内存连续，Triton 内核访问高效）
    assert key.stride(-1) == 1 and value.stride(-1) == 1  # 最后一维（head_dim）连续
    assert key.stride(1) == head_dim and value.stride(1) == head_dim  # 第 1 维（num_heads）步长为 head_dim
    assert k_cache.stride(1) == D and v_cache.stride(1) == D  # 缓存第 1 维步长为 D（单个元素总维度）
    assert slot_mapping.numel() == N  # slot 数量与 K/V 元素数一致

    # 3. 调用 Triton 内核：启动 N 个线程（每个线程处理 1 个 K/V 元素）
    store_kvcache_kernel[(N,)](  # 线程网格：(N,) 表示 N 个独立线程
        key, key.stride(0), value, value.stride(0),  # K/V 张量及第 0 维步长
        k_cache, v_cache, slot_mapping, D  # 缓存、slot 映射、维度 D
    )

class Attention(nn.Module):
    """
    3. Attention 类：注意力计算核心（支持 prefill/decode 双阶段）
注意力层是 LLM 的核心，需适配两个推理阶段：

Prefill 阶段：处理初始输入（如用户的提问文本），一次性计算所有输入 token 的 KV，并写入缓存。
Decode 阶段：生成后续 token（如模型的回答），每次仅计算 1 个新 token 的 Q，复用缓存中的 KV，避免重复计算。

关键参数解析：
num_kv_heads：支持 GQA（Grouped Query Attention） 或 MQA（Multi-Query Attention）。例如 num_heads=32、num_kv_heads=8，表示 32 个 Q 头共享 8 个 KV 头，大幅节省显存和计算量。
scale：对应之前分析的 “注意力得分缩放因子”，避免得分过大导致 softmax 梯度消失。
k_cache/v_cache：初始为空张量，推理时会预分配足够大的缓存（如 [max_seq_len, num_kv_heads, head_dim]）。
    """
    def __init__(
        self,
        num_heads: int,       # Query（Q）的头数（如 32）
        head_dim: int,        # 每个注意力头的维度（如 128）
        scale: float,         # 注意力得分缩放因子（如 1/sqrt(head_dim)）
        num_kv_heads: int,    # Key/Value（KV）的头数（支持 GQA/MQA，可小于 num_heads）
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = scale
        self.num_kv_heads = num_kv_heads
        # 初始化空 KV 缓存（推理时会根据实际需求分配大小）
        self.k_cache = self.v_cache = torch.tensor([])

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        """
         forward 方法：双阶段注意力计算
        """
        # 1. 重塑 QKV 形状：从 [batch*seq_len, num_heads*head_dim] 转为 [-1, num_heads, head_dim]
        # （-1 表示自动计算 batch*seq_len 的总元素数，适配批量处理）
        q = q.view(-1, self.num_heads, self.head_dim)
        k = k.view(-1, self.num_kv_heads, self.head_dim)
        v = v.view(-1, self.num_kv_heads, self.head_dim)

        # 2. 获取推理上下文（含阶段标识、slot 映射、块表等）
        context = get_context()  # 自定义上下文管理器，存储推理过程中的全局信息
        k_cache, v_cache = self.k_cache, self.v_cache

        # 3. 若 KV 缓存已初始化，将当前 K/V 写入缓存（Triton 加速）
        if k_cache.numel() and v_cache.numel():
            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)

        # 4. 分阶段计算注意力：Prefill 阶段（处理长序列）
        if context.is_prefill:
            # 4.1 若存在块表（block_tables），表示使用前缀缓存（prefix cache），直接复用缓存的 KV
            if context.block_tables is not None:
                k, v = k_cache, v_cache
            # 4.2 调用 FlashAttention 变长序列接口（支持不同长度的批量序列）
            o = flash_attn_varlen_func(
                q, k, v,
                max_seqlen_q=context.max_seqlen_q,  # Q 的最大序列长度
                cu_seqlens_q=context.cu_seqlens_q,  # Q 的累计序列长度（用于变长批量）
                max_seqlen_k=context.max_seqlen_k,  # K 的最大序列长度
                cu_seqlens_k=context.cu_seqlens_k,  # K 的累计序列长度
                softmax_scale=self.scale,  # 注意力缩放因子
                causal=True,  # 因果掩码（避免未来 token 影响当前，LLM 必需）
                block_table=context.block_tables  # KV 缓存块表（管理缓存块的复用）
            )

        # 5. 分阶段计算注意力：Decode 阶段（生成单个 token）
        else:
            # 调用 FlashAttention 带 KV 缓存的接口（复用历史 KV，仅计算新 Q 的注意力）
            o = flash_attn_with_kvcache(
                q.unsqueeze(1),  # Q 增加 batch 维度：[N, 1, num_heads, head_dim]
                k_cache, v_cache,  # 复用历史 KV 缓存
                cache_seqlens=context.context_lens,  # 每个序列的缓存长度
                block_table=context.block_tables,  # 缓存块表
                softmax_scale=self.scale,
                causal=True
            )

        # 6. 重塑输出：从 [-1, num_heads, head_dim] 转为 [-1, num_heads*head_dim]
        # （恢复与输入一致的特征维度，供后续 Feed-Forward 层使用）
        o = o.view(-1, self.num_heads * self.head_dim)
        return o
